feat(working-memory): measure working memory in admin calibration runs - #9897

@JPPhoto Thanks for the second pass. All six points are addressed in b012444. Each one has a test that fails without its fix.

High: shutdown overlapping old workers

• DefaultSessionProcessor.stop() now joins the workers against one shared 30 s deadline, after it signals stop and cancels their sessions. If it is still waiting after 1 s, it logs which worker and queue item it is waiting for.
• A worker still in its node at the deadline is logged as an error and recorded as abandoned. A repeated stop() does not wait for it again.
• While an abandoned worker is alive:
  • register_generation_devices() and start() refuse with a RuntimeError, so nothing replaces the device pool or starts a second worker on that GPU.
  • ApiDependencies.initialize refuses before it opens the database.
• The database is still disposed after the stop. A late transaction from an abandoned worker is refused by the database's thread claim.
• Why bounded: an unbounded join hangs shutdown forever on a node that never reaches a cancel point. The restart refusal closes the overlap a bounded join would otherwise leave.
• Tests: tests/app/services/test_invoker.py and tests/app/test_api_dependencies.py.

A detached locked record escaping the probe

• ModelCache.locks_held() now reads a counter that every lock and unlock keeps, independent of cache membership. That covers the lock/continue_lock failure paths too.
• CacheRecord.unlock() refuses an unlock without a lock before it decrements, so the counter cannot drift.
• Tests: test_model_cache_locks_held.py, including a dropped-while-locked record and both OOM paths.

reset_model undone by a pending write

• A reset raises a per-model epoch in the store file (model_reset_epochs), in the same file change, so the epoch exists only if the reset committed.
• The runner reads the epochs when the calibration node starts, before the probe scope opens.
• The write drops, under the file lock, every point attributed to a model whose epoch moved since then, including submodel windows such as <key>:vae.
• This also holds for a reset made by another process on the same root.
• Tests: a reset while the node runs, a reset by a second service instance, an in-flight write, and submodel attribution.

Retry lineage before admission

• The run's live read now returns retried_from_item_id (SessionQueueBase.get_queue_item_lineage, one indexed statement whose plan is pinned by a test). A retry stands in for its item before the runner starts it, and a canceled retry counts as the run's outcome.
• Remaining limit, documented: a retry that the queue prunes or clears before the runner ever saw it leaves no trace, so the run reads as the original item's outcome. Covering that would need a hook where the queue creates retries.

Keyed recovery after the model was removed

• When the run is recorded without items and the key's receipt exists, _prepare_run replays the receipt before any model or device lookup.
• A new request still requires the model.

Reset not invalidating the idempotency key

• A reset adds the removed keyed runs to removed_run_ids in the store file. Unkeyed runs are not added, since they cannot be repeated.
• A repeat of the keyed request after a reset is refused with the existing 404, whether the receipt is retained or already collected, on every device class.
• The reset run's pending items keep failing at admission.

Store file: the two new fields (removed_run_ids, model_reset_epochs) have empty defaults. The schema version stays 1, so existing files read unchanged.

Contract: no API shape or status changes. The OpenAPI descriptions are regenerated.

The stack above (#9899, #9910, #9919, #9927, #9946) has these merged in.

@JPPhoto Thanks. All five points are addressed in 1370b23, and each has a test that fails without its fix.

1. Status available on a store it cannot write. Status now checks that the current device class's store can be written. It opens an OS-deleted temporary file in the store folder, or in the nearest existing folder above it, and checks that the lock file is writable. On Windows it also checks that the JSON file is not read-only. A store that fails this reports available=false with the reason. device_class_ids() wraps listing errors as a store error, so an unreadable store directory answers 503 instead of 500. Tests cover a read-only lock file, a store root that is a file, and an unlistable directory.

2. A second batch after an ambiguous failure. A run now carries enqueue_attempted. A new run is recorded with it already set, in the same write. Only the attempt that set it clears it, and only when its enqueue is known to have queued nothing. A receipt replay never clears it. A repeat that finds a marked run with no items, no receipt and nothing queued or observed enqueues nothing. It returns the run with empty item_ids, and the run reads canceled. Tests: your sequence; a replay whose enqueue raises before the receipt is collected; and a full queue followed by a later repeat.

3. A completed item later reading canceled. On a clean finish (no error, not canceled, and the queue row is neither failed nor canceled), the runner now records the item as completed before the queue completes it. The final observation still writes the queue's status. Tests:

• the outcome survives a prune;
• the order is the runner's write, then the queue's;
• a failing node, and a failed save, are not recorded as completed.

5. 404 instead of 409. For a keyed request with no run, the reset check runs first (404), then the receipt check (409), and only then the model lookup.

7. Retry sources missing from item_ids. derive_run_status adds every lineage source to item_ids.

Limits, documented:

• If the process stops between the runner's completed record and the queue's completion, and the row is later pruned, the run reads completed. All of that item's nodes did finish.
• If clearing the mark fails, it is only logged. The run then stays marked, so that key cannot start it, though a request with another key can.

Store file: one new field with a default; schema version 1. API: no shape or status changes. The OpenAPI descriptions are regenerated.

The stack above has this merged in.

@JPPhoto Thanks. Both points are addressed in 9f2345c, with tests that fail without the fixes.

1. Cross-process creation

• A keyed run's creation now holds an OS file lock from preparing the run to recording its items. It covers the enqueue and the cleanup when nothing was queued.

• The lock uses the store's own lock primitive (msvcrt.locking / flock), so a holder that dies releases it with its process.

• It is striped over 256 files, working_memory/creations/<xx>.lock, where xx comes from the run id's sha256. That bounds the folder and means no lock file is ever deleted, which would race. The cost is false contention between unrelated runs that share a stripe, which is negligible for this admin action.

• A per-run in-process gate sits in front of it, so concurrent repeats in one process wait on the event loop rather than each holding an executor thread.

• A contender waits up to 30 s and then gets the existing retryable 503. After acquiring the lock it sees what the first attempt did:

  • the recorded items;
  • the receipt to replay;
  • or, once that receipt is collected, the marked run.

  So it can neither return an empty success while the first attempt is enqueuing, nor enqueue a second batch, nor replace the first attempt's item ids.

• Tests: two service instances on one root.

  • A repeat waits while the first instance is held inside its enqueue.
  • A second request cannot prepare while the first is recording the run.
  • Eight concurrent repeats in one process leave one worker on the file lock.

  Separate instances model separate processes here, because the OS arbitrates these locks per handle (LockFileEx) and per open file description (flock).

2. Final observation overwriting a saved completion

• When the queue no longer has the item, the final observation keeps a recorded completion.
• A queue row that still exists stays authoritative, including a failed one. So the run reads the same before and after a prune.
• Tests: completion recorded, item pruned, then the final observation; and a row that is failed in the queue.

Also changed: main's #9855 wake-up test. Merging main brought in #9855's test_worker_wakeup.py. Two of its places assumed that stop() returns while workers are still parked on test gates. Since stop() now joins its workers, the test (62d75d8):

• opens its gates before the cleanup stop;
• runs the stop case on a thread of its own, continuing once every worker's own notification is set.

#9855's guarantee is unchanged. Making the workers share one poll_now_event again fails all three cases.

The stack above has this merged in.

@JPPhoto Thanks, fixed in 9f87a27.

The fix

• The runner now passes queue_completed to the final observation: whether it saw the queue commit the item's completion, from the item _on_after_run_session returns.
• For an item the queue no longer has:
  • with queue_completed, the observation stays completed, so a completed item that is pruned keeps its outcome;
  • without it, it is canceled. That overrides a completion the runner recorded ahead of the queue, which covers a Clear Queue in that interval.
• A queue row that still exists keeps the last word, as before.
• I did not use the worker's cancel event. Clear Queue's handlers set it asynchronously, so it can arrive after the final observation.

Tests

• A runner test runs clear between the recorded completion and the queue's completion. The observation is canceled.
• A service test covers three cases: pruned after completion (completed), cleared before completion (canceled), and failed in the queue (failed).

Remaining limit (documented): if completing the item raises after the queue's commit, a prune or clear in the gap before the final observation records canceled. While the row exists it reads completed.

The stack above has this merged in.

Linux CUDA test run. Machine: Verda 2× NVIDIA RTX PRO 6000 Blackwell Server Edition (96 GB, sm_120), driver 580.178.04, Ubuntu 26.04.1, torch 2.13.0+cu130, cuDNN 9.20, Triton 3.7.1. The PR head was compared with its merge base on main; each GPU ran one lane, pinned by UUID.

Linux NVML reader, run on hardware

• pytest: tests/backend/util/test_nvml.py tests/backend/working_memory: 142 passed, including the slow hardware probe test.
• Probe scenario (SD VAE decode on this GPU):
  • The method is nvml and the driver delta is present.
  • The NVML per-process value matches nvidia-smi.
  • At 1024² the driver delta equals the reserved delta.
  • The first window (768²) reads +110 MiB over the reserved delta. That is first-use memory outside the allocator, gone in the next window.
• E2E start --calibrate through feat(e2e): calibrate working memory and build shipped data from probe runs #9899: 24 valid windows, all method nvml. scripts/working_memory/build_shipped.py reads the export.
  • Against the RTX 4090 Windows points under the same native allocator (shipped in vae.json), SD VAE fp32 decode measures 0.98× at 1024² and 0.99× at 1536², and the shipped estimate covers every Linux point (SD fp32/fp16, FLUX bf16 up to 3584²) with 6–10 % headroom. So the OS-independent (cudnn, native) keying holds here.
• Done in edba0dd: ("nvml", "cuda") is now in VERIFIED_METHODS, so new Linux NVIDIA windows say verified=true. The hardware test asserts that on Linux CUDA, and the docs note that containers without the host PID namespace read torch-only.

@JPPhoto Thanks — all four points are addressed in 49d8f8e, each with a test that fails without its fix.

1. Clear Queue, then a stop, left a false completed. I removed the early completion record (observe_completed) instead of making it provisional. Once the row is gone, a provisional record cannot be confirmed without a write on a status read, and reading it as canceled has the same outcome as having no early record. The trade-off: if the process stops after the queue completed an item and the item is then pruned before a restart, the item now reads canceled (the run partial if it measured valid points). Before, it read completed, so this reverses my round-3 fix for that window. It errs toward an honest outcome instead of a false success, and the measured points are kept. Test: clear just before the queue completes the item, the final observation lost, then a fresh service reads the run. It never reads completed.

2. A pruned failure lost failed and its error. The runner now hands the final observation the item as the queue settled it. It is captured right after the queue's completion or failure step, before statistics and callbacks run. When the row is gone the item keeps the outcome the queue gave it, completed or failed with its error. If the queue never settled it, the item reads failed when the runner raised and canceled otherwise. A row that still exists keeps the last word. Tests:

• a failure pruned before the observation;
• a runner error after Clear Queue;
• a session callback that raises after completion, followed by a prune.

3. NVML in containers without host PID visibility. On Linux CUDA the hardware test now asks NVML itself. If it finds no entry for this process, the test asserts the documented torch-only fallback. Where NVML matches, the full NVML assertions still apply.

4. The contradictory doc paragraph is rewritten. It now separates three cases: the queue settled the item before it was removed, Clear Queue removed it first, and the runner raised. It also states the stop-then-prune limit.

Also from @joshistoast's review on #9899: resets and the write fence now attribute <key>:<submodel> windows (e.g. model-a:vae) to their model. A model reset therefore removes them, and an in-flight write of them is fenced. The stored format is unchanged.

@JPPhoto Thanks, both are addressed in b80d534, each with a test that fails without its fix.

1. Terminal outcomes lost before observation. Every terminal queue transition in the runner now hands it the settled item at once. That covers the refusal, a revoked owner's cancel, a failed save of measurements, the node-error fail in _on_node_error, and the completion. If a transition raises after its commit, for example because the status announcement fails, the runner reads the row back and keeps it when it is terminal. A later prune therefore keeps the outcome the queue gave the item.

Tests: a refusal pruned before the final record, a node failure pruned before the session ended, and a completion whose announcement raised, then pruned.

Remaining limit: if the row is pruned in the moment between a raising announcement and the read-back, the item falls back to the runner's error and reads failed.

2. NVML in a private PID namespace. NVML attributes memory only when the process runs in the initial PID namespace, i.e. /proc/self/ns/pid has the init inode 0xEFFFFFFC. Elsewhere, including when the link cannot be read, the counters fall back to torch-only and the window is not marked verified. Bare-metal hosts and containers with --pid=host keep NVML.

I did not use NSpid: inside a container its first entry is the container's own id, not a host PID, so it would trust the same unknown mapping. The cost is that a container which shares the host's procfs but not its PID namespace also reads torch-only.

Test: a colliding PID in a namespace of its own reads unknown.